{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%html\n",
    "<style>\n",
    ".cell-output-ipywidget-background {\n",
    "    background-color: transparent !important;\n",
    "}\n",
    ":root {\n",
    "    --jp-widgets-color: var(--vscode-editor-foreground);\n",
    "    --jp-widgets-font-size: var(--vscode-editor-font-size);\n",
    "}  \n",
    "</style>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import json\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from openai.types.chat.chat_completion import ChatCompletion\n",
    "\n",
    "import art\n",
    "from art.local import LocalBackend\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "MODEL_NAME = \"001\"\n",
    "BASE_MODEL = \"Qwen/Qwen2.5-7B-Instruct\"\n",
    "TRAINING_STEPS = 1_000\n",
    "\n",
    "model = art.TrainableModel(\n",
    "    name=MODEL_NAME, project=\"rock-paper-tool-use\", base_model=BASE_MODEL\n",
    ")\n",
    "await model.register(LocalBackend())\n",
    "client = model.openai_client()\n",
    "\n",
    "\n",
    "def get_tool_call_id_and_move(chat_completion: ChatCompletion) -> tuple[str, str]:\n",
    "    tool_calls = chat_completion.choices[0].message.tool_calls\n",
    "    if not tool_calls:\n",
    "        return \"n/a\", \"nothing\"\n",
    "    tool_call = tool_calls[0]\n",
    "    try:\n",
    "        return tool_call.id, json.loads(tool_call.function.arguments)[\"move\"]\n",
    "    except json.JSONDecodeError:\n",
    "        return tool_call.id, \"nothing\"\n",
    "    except KeyError:\n",
    "        return tool_call.id, \"nothing\"\n",
    "\n",
    "\n",
    "async def rollout() -> art.Trajectory:\n",
    "    tools: art.Tools = [\n",
    "        {\n",
    "            \"type\": \"function\",\n",
    "            \"function\": {\n",
    "                \"name\": \"play_move\",\n",
    "                \"description\": \"Play a move in rock-paper-scissors\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"move\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"enum\": [\"rock\", \"paper\", \"scissors\"],\n",
    "                            \"description\": \"The move to play\",\n",
    "                        }\n",
    "                    },\n",
    "                    \"required\": [\"move\"],\n",
    "                },\n",
    "            },\n",
    "        }\n",
    "    ]\n",
    "    trajectories = [\n",
    "        art.Trajectory(\n",
    "            messages_and_choices=[\n",
    "                {\n",
    "                    \"role\": \"system\",\n",
    "                    \"content\": \"You are a rock-paper-scissors playing agent. Use the play_move function tool to declare your moves.\",\n",
    "                },\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": \"What will your first move be?\",\n",
    "                },\n",
    "            ],\n",
    "            tools=tools,\n",
    "            reward=0,\n",
    "            metrics={\n",
    "                \"num_rounds\": 0,\n",
    "                \"rock\": 0,\n",
    "                \"paper\": 0,\n",
    "                \"scissors\": 0,\n",
    "                \"nothing\": 0,\n",
    "            },\n",
    "        )\n",
    "        for _ in range(2)\n",
    "    ]\n",
    "    for _ in range(10):\n",
    "        chat_completions = await asyncio.gather(\n",
    "            *[\n",
    "                client.chat.completions.create(\n",
    "                    messages=trajectory.messages(),\n",
    "                    model=model,\n",
    "                    tools=tools,\n",
    "                    max_completion_tokens=100,\n",
    "                )\n",
    "                for trajectory, model in zip(trajectories, (MODEL_NAME, BASE_MODEL))\n",
    "            ]\n",
    "        )\n",
    "        for trajectory, chat_completion in zip(trajectories, chat_completions):\n",
    "            trajectory.messages_and_choices.append(chat_completion.choices[0])\n",
    "        (id0, move0), (id1, move1) = list(\n",
    "            map(get_tool_call_id_and_move, chat_completions)\n",
    "        )\n",
    "        beats = {\n",
    "            \"rock\": \"scissors\",\n",
    "            \"paper\": \"rock\",\n",
    "            \"scissors\": \"paper\",\n",
    "            \"nothing\": None,\n",
    "        }\n",
    "        if beats[move0] == move1:\n",
    "            trajectories[0].reward += 1\n",
    "        elif beats[move1] == move0:\n",
    "            trajectories[1].reward += 1\n",
    "        for trajectory in trajectories:\n",
    "            trajectory.metrics[\"num_rounds\"] += 1\n",
    "        trajectories[0].metrics[move0] += 1\n",
    "        trajectories[1].metrics[move1] += 1\n",
    "        if max(t.reward for t in trajectories) > 2:\n",
    "            break\n",
    "        trajectories[0].messages_and_choices.extend(\n",
    "            (\n",
    "                {\n",
    "                    \"role\": \"tool\",\n",
    "                    \"tool_call_id\": id0,\n",
    "                    \"content\": f\"The other player played {move1}.\",\n",
    "                },\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": \"What will your next move be?\",\n",
    "                },\n",
    "            )\n",
    "        )\n",
    "        trajectories[1].messages_and_choices.extend(\n",
    "            (\n",
    "                {\n",
    "                    \"role\": \"tool\",\n",
    "                    \"tool_call_id\": id1,\n",
    "                    \"content\": f\"The other player played {move0}.\",\n",
    "                },\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": \"What will your next move be?\",\n",
    "                },\n",
    "            )\n",
    "        )\n",
    "    return trajectories[0]\n",
    "\n",
    "\n",
    "for i in range(await model.get_step(), TRAINING_STEPS):\n",
    "    trajectories = await art.gather_trajectories(\n",
    "        (rollout() for _ in range(64)), max_exceptions=64\n",
    "    )\n",
    "    await model.train(\n",
    "        [art.TrajectoryGroup(trajectories)],\n",
    "        config=art.TrainConfig(learning_rate=5e-5),\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
